{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CSzhj5KSaFin"
   },
   "source": [
    "# Training Hello world model for Microcontrollers\n",
    "> In this post, we'll cover the basic tutorial for training simple regression model with tensorflow lite for for Microcontrollers(TFLM). This post is the summary of youtube video \"TinyML Book Screencast - Training the Hello World model\", presented by peter warden.\n",
    "\n",
    "- toc: true\n",
    "- badges: true\n",
    "- comments: true\n",
    "- author: Chanseok Kang\n",
    "- categories: [Python, Deep_Learning, Tensorflow-Keras, Tensorflow-Lite]\n",
    "- image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "id": "WHJ2M3AFaFio",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import tensorflow as tf\n",
    "# import numpy as np\n",
    "# import matplotlib.pyplot as plt\n",
    "# import pandas as pd\n",
    "\n",
    "# plt.rcParams['figure.figsize'] = (16, 10)\n",
    "# plt.rcParams['text.usetex'] = True\n",
    "# plt.rc('font', size=15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "roGsCLB-aFip"
   },
   "source": [
    "## Tensorflow Lite for Microcontrollers\n",
    "[TensorFlow Lite for Microcontrollers](https://www.tensorflow.org/lite/microcontrollers) (TFLM for short) is designed to run machine learning models on microcontrollers and other devices with only few kilobytes of memory. For the purpose of deploying machine learning model on embedded devices, it is also called [TinyML](https://tinyml.org/home/). As you know already, Tensorflow is one of commonly-used deep learning frameworks, and from version 2.x they offer some machine learning features for embedded systems via tensorflow lite. Unlike high performance mobile processor(like cortex-A series), Microcontroller (Cortex-M series or ESP32) has low power consumptions and it can deploy in various ways of customer products, like refrigerator, wash-machine and so on.\n",
    "\n",
    "Google introduced several supported boards for test,\n",
    "\n",
    "- [Arduino Nano 33 BLE Sense](https://store.arduino.cc/usa/nano-33-ble-sense-with-headers) (using Arduino IDE)\n",
    "- [SparkFun Edge](https://www.sparkfun.com/products/15170) (building directly from source)\n",
    "- [STM32F746 Discovery kit](https://www.st.com/en/evaluation-tools/32f746gdiscovery.html) (using Mbed)\n",
    "- [Adafruit EdgeBadge](https://www.adafruit.com/product/4400) (using Arduino IDE)\n",
    "- [Adafruit TensorFlow Lite for Microcontrollers Kit](https://www.adafruit.com/product/4317) (using Arduino IDE)\n",
    "- [Adafruit Circuit Playground Bluefruit](https://learn.adafruit.com/tensorflow-lite-for-circuit-playground-bluefruit-quickstart?view=all) (using Arduino IDE)\n",
    "- [Espressif ESP32-DevKitC](https://www.espressif.com/en/products/hardware/esp32-devkitc/overview) (using ESP IDF)\n",
    "- [Espressif ESP-EYE](https://www.espressif.com/en/products/hardware/esp-eye/overview) (using ESP IDF)\n",
    "\n",
    "Here, we'll implement the simple regression model in **Sparkfun Edge**. Most of contents are covered in Pete Warden's screencast. More informations are included in his [book](https://www.oreilly.com/library/view/tinyml/9781492052036/).\n",
    "\n",
    "> youtube: https://www.youtube.com/watch?v=8N6-WQsxwxA&t=205s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fmPYwadlaFip"
   },
   "source": [
    "### Prerequisites for implementation\n",
    "Before beginning, it requires some hardwares and basic knowledges,\n",
    "\n",
    "- hardware\n",
    "    - Sparkfun edge\n",
    "    - CH340c USB-c to UART converter\n",
    "    - Desktop PC for training model (or you can use google colab)\n",
    "- Knowledge\n",
    "    - Python\n",
    "    - C/C++\n",
    "    - Basic knowledge related on OS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "igRhnWpUaFiq"
   },
   "source": [
    "### Hello world\n",
    "Actually, \"Hello world\" may be the first program we faced, since it can show simple interaction between human and computer. TinyML also has simple example of \"Hello world\". Instead of Printing, we will build a model to generate the sine wave. So maybe our hypothesis will be like this,\n",
    "\n",
    "$$ \\tilde{H(x)} = \\sin(x) $$\n",
    "\n",
    "At first, we load the required packages,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yUmh-CjkaFiq"
   },
   "outputs": [],
   "source": [
    "# --encoding:gb2312--\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "\n",
    "\n",
    "gpus = tf.config.list_physical_devices('GPU')\n",
    "if gpus:\n",
    "    try:\n",
    "        # Currently, memory growth needs to be the same across GPUs\n",
    "        for gpu in gpus:\n",
    "            tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        logical_gpus = tf.config.list_logical_devices('GPU')\n",
    "        print(len(gpus), \"Physical GPUs,\", len(logical_gpus), \"Logical GPUs\")\n",
    "    except RuntimeError as e:\n",
    "        # Memory growth must be set before GPUs have been initialized\n",
    "        print(e)\n",
    "else:\n",
    "    print(\"No GPU was detected. CNNs can be very slow without a GPU.\")\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (16, 10)\n",
    "plt.rc('font', size=15)\n",
    "\n",
    "# define random seed for reproducibility\n",
    "np.random.seed(1) # numpy seed\n",
    "tf.random.set_seed(1) # tensorflow global random seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PCxkga35aFiq",
    "outputId": "278ac161-4157-4b3b-d7bd-a529534afc23"
   },
   "outputs": [],
   "source": [
    "print('Numpy: {}'.format(np.__version__))\n",
    "print('Tensorflow: {}'.format(tf.__version__))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_9qfC4O9aFir"
   },
   "source": [
    "To train the model, it requires a sort of data, namely training data. In our case, we will sample the random data from numpy and generate the actual output from known model(the sine function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 426
    },
    "id": "OEPibVFiaFir",
    "outputId": "f612f9e9-028b-44ca-9e69-24903eef984e"
   },
   "outputs": [],
   "source": [
    "# Generate a uniformly distributed set of random numbers in the range from\n",
    "# 0 to 2π, which covers a complete sine wave oscillation\n",
    "X = np.random.uniform(\n",
    "    low=0, high=2*math.pi, size=10000).astype(np.float32)\n",
    "\n",
    "# Shuffle the values to guarantee they're not in order\n",
    "np.random.shuffle(X)\n",
    "\n",
    "# Calculate the corresponding sine values\n",
    "y = np.sin(X).astype(np.float32)\n",
    "\n",
    "# Plot our data. The 'b.' argument tells the library to print blue dots.\n",
    "plt.plot(X, y, 'b.')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6wE5odM3aFir"
   },
   "source": [
    "But this data doesn't reflect the real-world data, because there is no variation of distribution in dataset, also known as noise. We can add random data for each x, so it makes to seem more randomly distributed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 426
    },
    "id": "sYTtIMhyaFis",
    "outputId": "4b4130b6-ea5e-4cec-9633-67eed36f62bb"
   },
   "outputs": [],
   "source": [
    "# Add a small random number to each y value\n",
    "y += 0.1 * np.random.randn(*y.shape)\n",
    "\n",
    "# Plot our data\n",
    "plt.plot(X, y, 'b.')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UpoGm4IXaFis"
   },
   "source": [
    "Through this, we expect that the model is approximated with sinusoid curve if it is well trained.\n",
    "\n",
    "To use it for training, we are going to preprocess the data. We'll split it with following proportions,\n",
    "    \n",
    "- Train data: 60%\n",
    "- Validation data: 20%\n",
    "- Test data: 20%"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 426
    },
    "id": "skH4e1ePaFis",
    "outputId": "7ad4cfdf-7d39-470e-9049-4dfa1fcf4221"
   },
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=1)\n",
    "X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.25, random_state=1)\n",
    "\n",
    "# Plot the data in each partition in different colors:\n",
    "plt.plot(X_train, y_train, 'b.', label=\"Train\")\n",
    "plt.plot(X_val, y_val, 'y.', label=\"Validate\")\n",
    "plt.plot(X_test, y_test, 'r.', label=\"Test\")\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dR7GVCD1aFis"
   },
   "source": [
    "To train the model, we will use sequential model in tensorflow-keras, and add two Dense layer. Then we will add adam optimizer, and set mean squared error for loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 248
    },
    "id": "pFYX5vQDaFis",
    "outputId": "23aee027-0338-40e2-c727-37d22c226189"
   },
   "outputs": [],
   "source": [
    "model = tf.keras.Sequential(name='sine')\n",
    "model.add(tf.keras.layers.Dense(16, activation='relu', input_shape=(1, )))\n",
    "model.add(tf.keras.layers.Dense(1))\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8ocihZhSaFis"
   },
   "outputs": [],
   "source": [
    "model.compile(optimizer='adam', loss='mse', metrics=['accuracy', 'mae'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aWC8I5d7aFis"
   },
   "source": [
    "Now, it's time to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4xOfTHwiaFis"
   },
   "outputs": [],
   "source": [
    "history = model.fit(X_train, y_train, epochs=1000, batch_size=15, validation_data=(X_val, y_val), verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tE90kCpjaFis"
   },
   "outputs": [],
   "source": [
    "# Draw a graph of the loss, which is the distance between\n",
    "# the predicted and actual values during training and validation.\n",
    "loss = history.history['loss']\n",
    "val_loss = history.history['val_loss']\n",
    "\n",
    "epochs = range(1, len(loss) + 1)\n",
    "\n",
    "plt.plot(epochs, loss, 'g.', label='Training loss')\n",
    "plt.plot(epochs, val_loss, 'b', label='Validation loss')\n",
    "plt.title('Training and validation loss')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eiFDl33naFit"
   },
   "outputs": [],
   "source": [
    "# Exclude the first few epochs so the graph is easier to read\n",
    "SKIP = 500\n",
    "\n",
    "plt.plot(epochs[SKIP:], loss[SKIP:], 'g.', label='Training loss')\n",
    "plt.plot(epochs[SKIP:], val_loss[SKIP:], 'b.', label='Validation loss')\n",
    "plt.title('Training and validation loss')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mncv80zjaFit"
   },
   "outputs": [],
   "source": [
    "# Calculate and print the loss on our test dataset\n",
    "loss = model.evaluate(X_test, y_test)\n",
    "\n",
    "# Make predictions based on our test dataset\n",
    "predictions = model.predict(X_test)\n",
    "\n",
    "# Graph the predictions against the actual values\n",
    "plt.clf()\n",
    "plt.title('Comparison of predictions and actual values')\n",
    "plt.plot(X_test, y_test, 'b.', label='Actual')\n",
    "plt.plot(X_test, predictions, 'r.', label='Predicted')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IH-1L5U_aFit"
   },
   "outputs": [],
   "source": [
    "model_2 = tf.keras.Sequential(name='sine2')\n",
    "\n",
    "# First layer takes a scalar input and feeds it through 16 \"neurons\". The\n",
    "# neurons decide whether to activate based on the 'relu' activation function.\n",
    "model_2.add(tf.keras.layers.Dense(16, activation='relu', input_shape=(1,)))\n",
    "\n",
    "# The new second layer may help the network learn more complex representations\n",
    "model_2.add(tf.keras.layers.Dense(16, activation='relu'))\n",
    "\n",
    "# Final layer is a single neuron, since we want to output a single value\n",
    "model_2.add(tf.keras.layers.Dense(1))\n",
    "\n",
    "# Compile the model using a standard optimizer and loss function for regression\n",
    "model_2.compile(optimizer='adam', loss='mse', metrics=['accuracy', 'mae'])\n",
    "\n",
    "model_2.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fjwqmBb-aFit"
   },
   "outputs": [],
   "source": [
    "history = model_2.fit(X_train, y_train, epochs=500, batch_size=64,\n",
    "                    validation_data=(X_val, y_val), verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QDeJ-n1zaFit"
   },
   "outputs": [],
   "source": [
    "# Draw a graph of the loss, which is the distance between\n",
    "# the predicted and actual values during training and validation.\n",
    "loss = history.history['loss']\n",
    "val_loss = history.history['val_loss']\n",
    "\n",
    "epochs = range(1, len(loss) + 1)\n",
    "\n",
    "# Exclude the first few epochs so the graph is easier to read\n",
    "SKIP = 100\n",
    "\n",
    "plt.figure()\n",
    "plt.subplot(1, 2, 1)\n",
    "\n",
    "plt.plot(epochs[SKIP:], loss[SKIP:], 'g.', label='Training loss')\n",
    "plt.plot(epochs[SKIP:], val_loss[SKIP:], 'b.', label='Validation loss')\n",
    "plt.title('Training and validation loss')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "\n",
    "# Draw a graph of mean absolute error, which is another way of\n",
    "# measuring the amount of error in the prediction.\n",
    "mae = history.history['mae']\n",
    "val_mae = history.history['val_mae']\n",
    "\n",
    "plt.plot(epochs[SKIP:], mae[SKIP:], 'g.', label='Training MAE')\n",
    "plt.plot(epochs[SKIP:], val_mae[SKIP:], 'b.', label='Validation MAE')\n",
    "plt.title('Training and validation mean absolute error')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('MAE')\n",
    "plt.legend()\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x-OkwXO6aFit"
   },
   "outputs": [],
   "source": [
    "# Calculate and print the loss on our test dataset\n",
    "loss = model_2.evaluate(X_test, y_test)\n",
    "\n",
    "# Make predictions based on our test dataset\n",
    "predictions = model_2.predict(X_test)\n",
    "\n",
    "# Graph the predictions against the actual values\n",
    "plt.clf()\n",
    "plt.title('Comparison of predictions and actual values')\n",
    "plt.plot(X_test, y_test, 'b.', label='Actual')\n",
    "plt.plot(X_test, predictions, 'r.', label='Predicted')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JbfCAh2kaFit"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "MODEL_DIR = './models/'\n",
    "if not os.path.exists(MODEL_DIR):\n",
    "    os.mkdir(MODEL_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sE1lITdraFit"
   },
   "outputs": [],
   "source": [
    "# Convert the model to the TensorFlow Lite format without quantization\n",
    "converter = tf.lite.TFLiteConverter.from_keras_model(model_2)\n",
    "model_no_quant_tflite = converter.convert()\n",
    "\n",
    "# # Save the model to disk\n",
    "open(MODEL_DIR + 'model_no_quant.tflite', \"wb\").write(model_no_quant_tflite)\n",
    "\n",
    "# Convert the model to the TensorFlow Lite format with quantization\n",
    "def representative_dataset():\n",
    "    for i in range(500):\n",
    "        yield([X_train[i].reshape(1, 1)])\n",
    "# Set the optimization flag.\n",
    "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "# Enforce full-int8 quantization (except inputs/outputs which are always float)\n",
    "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
    "# Provide a representative dataset to ensure we quantize correctly.\n",
    "converter.representative_dataset = representative_dataset\n",
    "model_tflite = converter.convert()\n",
    "\n",
    "# Save the model to disk\n",
    "open(MODEL_DIR + 'model.tflite', \"wb\").write(model_tflite)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BD7Yq4qRaFit"
   },
   "outputs": [],
   "source": [
    "model_no_quant_size = os.path.getsize(MODEL_DIR + 'model_no_quant.tflite')\n",
    "print(\"Model is %d bytes\" % model_no_quant_size)\n",
    "model_size = os.path.getsize(MODEL_DIR + 'model.tflite')\n",
    "print(\"Quantized model is %d bytes\" % model_size)\n",
    "difference = model_no_quant_size - model_size\n",
    "print(\"Difference is %d bytes\" % difference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KOG3alZZaFiu"
   },
   "outputs": [],
   "source": [
    "# Instantiate an interpreter for each model\n",
    "model_no_quant = tf.lite.Interpreter(MODEL_DIR + 'model_no_quant.tflite')\n",
    "model = tf.lite.Interpreter(MODEL_DIR + 'model.tflite')\n",
    "\n",
    "# Allocate memory for each model\n",
    "model_no_quant.allocate_tensors()\n",
    "model.allocate_tensors()\n",
    "\n",
    "# Get the input and output tensors so we can feed in values and get the results\n",
    "model_no_quant_input = model_no_quant.tensor(model_no_quant.get_input_details()[0][\"index\"])\n",
    "model_no_quant_output = model_no_quant.tensor(model_no_quant.get_output_details()[0][\"index\"])\n",
    "model_input = model.tensor(model.get_input_details()[0][\"index\"])\n",
    "model_output = model.tensor(model.get_output_details()[0][\"index\"])\n",
    "\n",
    "# Create arrays to store the results\n",
    "model_no_quant_predictions = np.empty(X_test.size)\n",
    "model_predictions = np.empty(X_test.size)\n",
    "\n",
    "# Run each model's interpreter for each value and store the results in arrays\n",
    "for i in range(X_test.size):\n",
    "    model_no_quant_input().fill(X_test[i])\n",
    "    model_no_quant.invoke()\n",
    "    model_no_quant_predictions[i] = model_no_quant_output()[0]\n",
    "\n",
    "    model_input().fill(X_test[i])\n",
    "    model.invoke()\n",
    "    model_predictions[i] = model_output()[0]\n",
    "\n",
    "# See how they line up with the data\n",
    "plt.clf()\n",
    "plt.title('Comparison of various models against actual values')\n",
    "plt.plot(X_test, y_test, 'bo', label='Actual values', alpha=0.4)\n",
    "plt.plot(X_test, predictions, 'ro', label='Original predictions')\n",
    "plt.plot(X_test, model_no_quant_predictions, 'bx', label='Lite predictions')\n",
    "plt.plot(X_test, model_predictions, 'gx', label='Lite quantized predictions')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ITqmSbnCaFiu"
   },
   "outputs": [],
   "source": [
    "!xxd -i {MODEL_DIR + 'model.tflite'} > {MODEL_DIR + 'model.cc'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Pmcjml1taFiu"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VbvkKSiuaFiu"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "mytf_gpu",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
